# -*- coding:utf-8 -*-
'''
根据字符批量生成验证码：有正负，3，4，5 位
由于考虑随机性太麻烦，为简单分开考虑
命名：序号_位数_编号.jpg
屏幕一共可显示 7 位
'''
import os
from PIL import Image,ImageFilter
import numpy as np
import random  # 在文件顶部添加

# 获取当前脚本所在目录的绝对路径
script_dir = os.path.dirname(os.path.abspath(__file__))
# 基于脚本目录构建save_path
save_path = os.path.join(script_dir, '../gen_digit/')
if not os.path.exists(save_path):
  os.mkdir(save_path)

# 字符路径
char_path = os.path.join(script_dir,'../digits/')
print('字符路径：',char_path)
# 获取目录中第一个图片文件的后缀（更可靠）
image_files = [f for f in os.listdir(char_path) if f.lower().endswith(('.png', '.jpg', '.jpeg', '.bmp'))]
if not image_files:
    raise FileNotFoundError(f"在 {char_path} 中未找到图片文件")
houz = os.path.splitext(image_files[0])[1]
pj_size = Image.open(char_path + '0' + houz).size
print('图片尺寸：',pj_size)
new_size = 7*pj_size[0],pj_size[1]
required_files = ['bl', 'bc', 'fc', 'fl', 'br', 'dot']  # 添加所有必要的基础文件名
for file in required_files:
    file_path = os.path.join(char_path, f'{file}{houz}' if file != 'dot' else f'{file}.jpg')
    if not os.path.exists(file_path):
        raise FileNotFoundError(f"必要文件缺失: {file_path}")

def gen_3(num,target):
  # 生成 3 位的随机数
  name = [str(num) + '_' + '3_']
  img = Image.open(char_path + 'bl' + houz)  # blank left
  target.paste(img, (0, 0, pj_size[0], pj_size[1]))# 将image复制到target的指定位置中

  img = Image.open(char_path + 'bc' + houz)  # blank center
  target.paste(img, (pj_size[0], 0, 2*pj_size[0], pj_size[1]))# 将image复制到target的指定位置中

  if '-' == random.choice(['+', '-']):  # 使用 Python 原生 random.choice
    name.append('-')
    img = Image.open(char_path + 'fc' + houz)
    target.paste(img, (2*pj_size[0], 0, 3*pj_size[0], pj_size[1]))# 将image复制到target的指定位置中
  else:
    img = Image.open(char_path + 'bc' + houz)
    target.paste(img, (2*pj_size[0], 0, 3*pj_size[0], pj_size[1]))# 将image复制到target的指定位置中

  for i in range(3):
    num = str(np.random.choice(range(10)))
    name.append(num)
    img = Image.open(char_path + num + houz)
    target.paste(img,box=(pj_size[0]*(i+3), 0, pj_size[0]*(i+4), pj_size[1]))
  return name,target

def gen_4(num,target):
  # 生成 4 位的随机数
  name = [str(num) + '_' + '4_']
  img = Image.open(char_path + 'bl' + houz)  # blank left
  target.paste(img, (0, 0, pj_size[0], pj_size[1]))# 将image复制到target的指定位置中

  if '-' == np.random.choice(['+','-']):
    name.append('-')
    img = Image.open(char_path + 'fc' + houz)
    target.paste(img, (1*pj_size[0], 0, 2*pj_size[0], pj_size[1]))# 将image复制到target的指定位置中
  else:
    img = Image.open(char_path + 'bc' + houz)
    target.paste(img, (1*pj_size[0], 0, 2*pj_size[0], pj_size[1]))# 将image复制到target的指定位置中

  for i in range(4):
    num = str(np.random.choice(range(10)))
    name.append(num)
    img = Image.open(char_path + num + houz)
    target.paste(img,box=(pj_size[0]*(i+2), 0, pj_size[0]*(i+3), pj_size[1]))
  return name,target

def gen_5(num,target):
  # 生成 5 位的随机数
  name = [str(num) + '_' + '5_']

  if '-' == np.random.choice(['+','-']):
    name.append('-')
    img = Image.open(char_path + 'fl' + houz)
    target.paste(img, (0, 0, pj_size[0], pj_size[1]))# 将image复制到target的指定位置中
  else:
    img = Image.open(char_path + 'bl' + houz)
    target.paste(img, (0, 0, pj_size[0], pj_size[1]))# 将image复制到target的指定位置中

  for i in range(5):
    num = str(np.random.choice(range(10)))
    name.append(num)
    img = Image.open(char_path + num + houz)
    target.paste(img,box=(pj_size[0]*(i+1), 0, pj_size[0]*(i+2), pj_size[1]))
  return name,target

def gen_pic(pic_num):
    try:
        # 现有gen_pic逻辑
        ty = np.random.choice([3,4,5])
        # ty = 3
        target = Image.new('RGB', new_size)
        img = Image.open(char_path + 'br' + houz)  # blank right
        target.paste(img, (6*pj_size[0], 0, new_size[0], new_size[1]))# 将image复制到target的指定位置中
        if ty == 3:
            name,target = gen_3(pic_num,target)
        elif ty == 4:
            name,target = gen_4(pic_num,target)
        else:
            name,target = gen_5(pic_num,target)
        quality_value = 50 # quality来指定生成图片的质量，范围是0～100

        # 相对于图片右下角的动态坐标
        dot_position = (new_size[0] - 50, new_size[1] - 50)  # 距离右下角50px
        img = Image.open(os.path.join(char_path, 'dot.jpg'))
        target.paste(img, dot_position)
        target.paste(img, (940, 550))
        # # 对图片进行变换
        # rank = np.random.choice([0.3,0.6,1,2,3,4,5])
        # target = target.point(lambda p:p*rank)   # 对图片亮度进行随机调节 np.random.choice([0.5,1,2,3,4])
        # # target = target.point(lambda p:p*np.random.choice([0.5,1,2,3,4])) # 将rank 放入可生成彩色图片
        # target = target.filter(ImageFilter.GaussianBlur(radius=np.random.choice(range(2,7)))) # 对图片在 2-7 进行随机模糊化
        # target = target.rotate(np.random.choice(range(-20,20))) # 对图片在 20 度内进行随机旋转
        # target = target.resize((160,60))#.convert("L")
        pic_path = save_path + ''.join(name) + houz
        target.save(pic_path, quality = quality_value)
    except Exception as e:
        print(f"生成图片 {pic_num} 失败: {str(e)}")
        return None

def main():
  for i in range(10):
    gen_pic(i)

if __name__ == '__main__':
  # 图片保存路径
  main()